package scu.maqiang.elasticity;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES2T32;
import scu.maqiang.mesh.*;
import scu.maqiang.numeric.*;

public class Elasticity2DT32 {
	
	public Elasticity2DT32(Mesh2T3 mesh) {
		this.mesh = mesh;
		fs = new FES2T32(mesh);
		K = new SRMatrix(fs.GetNdof());
		x = new double[fs.GetNdof()];
		RHS = new double[fs.GetNdof()];
		disp = new double[2][mesh.getNv()];
		strain = new double[4][mesh.getNv()];
		stress = new double[4][mesh.getNv()];
		eqStress = new double[mesh.getNv()];
	}

	public static void main(String[] args) {
		long startTime = System.nanoTime();
		/*
		Mesh2DT3Generator mg = new Mesh2DT3Generator();
		Mesh2DT3 mesh = mg.square2D(100, 10, T3Type.Right, Direct.Y);
		mesh.scale(1.0, 0.1);
		Elasticity2DT32 e2dt32 = new Elasticity2DT32(mesh);
		ScalarCoefFunc EFunc = (xy, label, param) -> 100.0;
		ScalarCoefFunc NuFunc = (xy, label, param) -> 0.3;
		ScalarCoefFunc vfu = (xy, label, param) -> 0.0;
		ScalarCoefFunc vfv = (xy, label, param) -> -10.0;
		ScalarCoefFunc dispU = (xy, label, param) -> 0.0;
		ScalarCoefFunc dispV = (xy, label, param) -> 0.0;
		e2dt32.elastoStaticProblem(BVPType.PLANESTRESS, new ScalarCoefFunc[] {EFunc, NuFunc}, null, 
				new ScalarCoefFunc[] {vfu, vfv}, null, 
				new int[][] {new int[] {4}}, new Direct[] {Direct.All}, new ScalarCoefFunc[][] {new ScalarCoefFunc[] {dispU, dispV}}, 
				new double[][] {new double[] {0.0}}, 
				null, null, null, "MyPlaneStress2DT32.dat");
		*/

		Mesh2T3 mesh = new Mesh2T3().square2D(100, 10, T3Type.Cross2, Direct.X);
		mesh.scale(10.0, 1.0);
		Elasticity2DT32 e2dt32 = new Elasticity2DT32(mesh);
		e2dt32.elastoStaticProblem(BVPType.PLANESTRAIN, new double[] {2.0, 0.3}, new double[]{0, -1.0},
				new int[][] {new int[] {4}}, new Direct[] {Direct.All}, new double[][] {new double[] {0.0, 0.0}}, 
				null, null,
				null, null, null, "MyPlaneStress2DT323.dat");	
		

		//		Mesh2DT3Generator mg = new Mesh2DT3Generator();
//		Mesh2DT3 mesh = mg.square2D(100, 10, T3Type.Right, Direct.Y);
//		mesh.scale(1.0, 0.1);
//		FES2T32 fs = new FES2T32(mesh);
//		SRMatrix A = new SRMatrix(fs.GetNdof());
//		fs.assembleElasticStiff(new double[] {100.0, 0.3}, BVPType.PLANESTRESS, A);
//		//A.sort();
//		//System.out.println(A);
//        System.out.println("Size of A: " + A.getNNZ());
//        double[] RHS = new double[fs.GetNdof()];
//        fs.assembleVolumeForce(new double[] {0.0, -10.0}, BVPType.COMMON, RHS);
//        
//        fs.applyBCZero(A, RHS, Direct.All, BCType.MBN, 4);
////        System.out.println("RHS = " + MVector.toString(RHS));
//        IterSSolver solver = new IterSSolver(1e-10, 3000);
//        double[] x = new double[fs.GetNdof()];
//        
//        solver.PCGSSOR(A, RHS, x, 1.7, 1);
//        
//        double[][] uv = new double[2][mesh.getNv()];
//        fs.extractComponent(x, uv[0], uv[1], null);
//        mesh.toTecplot("PlaneStress2DT32.dat", uv);
//        
//        NewIterSSolver newSolver = new NewIterSSolver(A);
//        Arrays.fill(x, 0.0);
//        newSolver.GMRES_SSOR(RHS, x, 1.5, 100, 2);
//        
//        Arrays.fill(x, 0.0);
//        newSolver.PCGSSOR(RHS, x, 1.5, 1);
		
//		Mesh2DT3Generator mg = new Mesh2DT3Generator();
//		Mesh2DT3 mesh = mg.square2D(80, 10, T3Type.Left);
//		mesh.scale(8.0, 1.0);
//		Elasticity2DT32 e2dt32 = new Elasticity2DT32(mesh);
//		e2dt32.elastoStaticProblem(BVPType.PLANESTRESS, new double[] {100, 0.3}, new double[] {0.0, -10.0}, 
//				new int[][] {new int[] {4}}, new Direct[] {Direct.All}, new double[][] {new double[] {0.0, 0.0}}, 
//				null, null, "MyPlaneStress2DT323.dat");		
        
        long endTime = System.nanoTime();
        System.out.println("Elapse Time: " + (endTime - startTime)/1e9 + " s");
	}
	
	
	/**
	 * 
	 * @param tp
	 * @param elasticCoefFunc
	 * @param paramElasticCoef
	 * @param volumeForce
	 * @param paramVolumeForce
	 * @param fixedDispLabel
	 * @param direct
	 * @param fixedDisp
	 * @param tractionLabel
	 * @param traction
	 * @param paramTraction
	 * @param reportFile
	 */
	public void elastoStaticProblem(BVPType tp, ScalarFunc[] elasticCoefFunc, double[] paramElasticCoef,
			                       ScalarFunc[] volumeForce, double[] paramVolumeForce,
			                       int[][] fixedDispLabel, Direct[] direct, ScalarFunc[][] fixedDisp, double[][] paramFixedDisp, 
			                       int[] tractionLabel, ScalarFunc[] traction, double[] paramTraction, String reportFile) {
		fs.assembleStiff(elasticCoefFunc, paramElasticCoef, tp, K);
		if (volumeForce != null) {
			fs.assembleVolumeForce(volumeForce, paramVolumeForce, tp, RHS);
		}
		//System.out.println(MVO.toString(RHS));
		
		if (tractionLabel != null) {
			fs.assembleSurfaceForce(traction, paramTraction, tp, RHS, tractionLabel);
		}
		
		int n = direct.length;
		ParamCheck.checkEqual(fixedDispLabel.length, direct.length);
		for(int i = 0; i < n; i++) {
			fs.applyBC_MBN(K, direct[i], fixedDispLabel[i]);
		}
		
		if(fixedDisp != null) {
			for(int i = 0; i < n; i++) {
				fs.applyBC_MBN(RHS, direct[i], fixedDisp[i], paramFixedDisp[i], fixedDispLabel[i]);
			}
		} else {
			for(int i = 0; i < n; i++) {
				fs.applyBC_MBN(RHS, direct[i], 0.0, fixedDispLabel[i]);
			}
		}
		
		//System.out.println(MVO.toString(RHS));
		//solveAndComputeStrainStressAndExport(reportFile);
		
	}
	
	public void elastoStaticProblem(BVPType tp, double[] elasticCoef, double[] volumeForce,
			                        int[][] fixedDispLabel, Direct[] direct, double[][] fixedDisp, 
			                        int[][] tractionLabel, double[][] traction,  
			                        int[][] pressureLabel, double[] pressure, 
			                        double[][] heatIncrement,
			                        String reportFile) {
		fs.assembleStiff(elasticCoef, tp, K);
		if (volumeForce != null) {
			fs.assembleVolumeForce(volumeForce, tp, RHS);
		}

		if (tractionLabel != null) {
			int n = tractionLabel.length;
			ParamCheck.checkEqual(n, traction.length);
			for(int i = 0; i < n; i++) {
				fs.assembleSurfaceForce(traction[i], tp, RHS, tractionLabel[i]);
			}
		}
		
		if (pressureLabel != null) {
			int n = pressureLabel.length;
			ParamCheck.checkEqual(n, pressure.length);
			for(int i = 0; i < n; i++) {
				fs.assemblePressure(new double[] {pressure[i]}, tp, RHS, tractionLabel[i]);
			}
		}
		
		if (heatIncrement != null) {
			fs.assembleHeatLoad(new double[] {elasticCoef[2]}, heatIncrement, tp, RHS);
		}
		
		//K.sort();
		//System.out.println(K.display());
		//System.out.println(MVO.toString(RHS));
		
		int n = direct.length;
		ParamCheck.checkEqual(fixedDispLabel.length, direct.length);
		for(int i = 0; i < n; i++) {
			fs.applyBC_MBN(K, direct[i], fixedDispLabel[i]);
		}
		
		if(fixedDisp != null) {
			for(int i = 0; i < n; i++) {
				fs.applyBC_MBN_Disp(RHS, direct[i], fixedDisp[i], fixedDispLabel[i]);
			}
		} else {
			for(int i = 0; i < n; i++) {
				fs.applyBC_MBN(RHS, direct[i], 0.0, fixedDispLabel[i]);
			}
		}
		
		//System.out.println(MVO.toString(RHS));
		solveAndComputeStrainStressAndExport(tp, elasticCoef, reportFile);
	}
	
	public void solveAndComputeStrainStressAndExport(BVPType tp, double[] elasticCoef, String reporFile) {
		//NewIterSSolver solver = new NewIterSSolver(K, false);
		//solver.setEps(1.0e-16);
		//solver.PCGSSOR(RHS, x, 1.50, 1);
		DirectSSolver solver = new DirectSSolver(K);
		x = solver.CholeskySolve(RHS, 1);
		//DirectSSolver dsolver = new DirectSSolver(K, fs.getnBand());
		//dsolver.PLUSolve(RHS, x, 1);
		fs.extractComponent(x, disp);
		fs.computeNodalStrainStressVonMisesStress(tp, elasticCoef, x, strain, stress, eqStress);
		
		System.out.println("Displacement: ");
		System.out.println("u min: " + MVO.min(disp[0]) + "\tmax: " + MVO.max(disp[0]));
		System.out.println("v min: " + MVO.min(disp[1]) + "\tmax: " + MVO.max(disp[1]));
		System.out.println("Strain: ");
		System.out.println("ex  min: " + MVO.min(strain[0]) + "\tmax: " + MVO.max(strain[0]));
		System.out.println("ey  min: " + MVO.min(strain[1]) + "\tmax: " + MVO.max(strain[1]));
		System.out.println("exy min: " + MVO.min(strain[2]) + "\tmax: " + MVO.max(strain[2]));
		System.out.println("ez  min: " + MVO.min(strain[3]) + "\tmax: " + MVO.max(strain[3]));
		System.out.println("Stress: ");
		System.out.println("tx  min: " + MVO.min(stress[0]) + "\tmax: " + MVO.max(stress[0]));
		System.out.println("ty  min: " + MVO.min(stress[1]) + "\tmax: " + MVO.max(stress[1]));
		System.out.println("txy min: " + MVO.min(stress[2]) + "\tmax: " + MVO.max(stress[2]));
		System.out.println("tz  min: " + MVO.min(stress[3]) + "\tmax: " + MVO.max(stress[3]));
		
		System.out.println("Equvilent Stress: ");
		System.out.println("tx  min: " + MVO.min(eqStress) + "\tmax: " + MVO.max(eqStress));
//		fs.computeNodalStrain(x, strain);
//		fs.computeNodalStress(x, stress);
//		fs.computeEqStress(x, eqStress);
//		fs.computeNodalStrainStress(x, strain,stress, eqStress);
		double[][] resultVector = new double[11][];
		resultVector[0] = disp[0];
		resultVector[1] = disp[1];
		resultVector[2] = strain[0];
		resultVector[3] = strain[1];
		resultVector[4] = strain[2];
		resultVector[5] = strain[3];
		resultVector[6] = stress[0];
		resultVector[7] = stress[1];
		resultVector[8] = stress[2];
		resultVector[9] = stress[3];
		resultVector[10] = eqStress;
		mesh.toTecplot(reporFile, resultVector);
	}
	
	public double[][] exportDisplacement() {
		double[][] myDisp = new double[2][mesh.getNv()];
		myDisp[0] = disp[0].clone();
		myDisp[1] = disp[1].clone();
		return myDisp;
	}
	
	Mesh2T3 mesh;
	FES2T32 fs;
	SRMatrix K;
	double[] RHS;
	double[] x;
	double[][] disp;
	double[][] strain;
	double[][] stress;
	double[] eqStress;
}
